import cv2
import face_recognition
import numpy as np


def save_face_reco_handle(num, filename, facename):
    # 录入人脸库
    try:
        global KNOW_FACE_DIR
        global KNOW_FACE_ENCODING
        global KNOW_FACE_NAMES
    except:
        print("请在主程序中导入人脸识别初始化模块")
        return None

    WIDTH = 240
    HEIGHT = 320
    print(f"开始录入{facename}的人脸照片...")
    try:
        cap = cv2.VideoCapture(num)
        while not cap.isOpened():
            continue
        # cv2.waitKey(1)
        cvimg_success, img_src = cap.read()
        cvimg_h, cvimg_w, cvimg_c = img_src.shape
        cvimg_w1 = cvimg_h*WIDTH//HEIGHT
        cvimg_x1 = (cvimg_w-cvimg_w1)//2
        img_src = img_src[:, cvimg_x1:cvimg_x1+cvimg_w1]
        img_src = cv2.resize(img_src, (WIDTH, HEIGHT))
        known_faces_dir = f"./{KNOW_FACE_DIR}/{filenam}"
        print(known_faces_dir)
        cv2.imwrite(known_faces_dir, img_src)

        image = face_recognition.load_image_file(known_faces_dir)
        encoding = face_recognition.face_encodings(image)[0]
        KNOW_FACE_ENCODING.append(encoding)
        KNOW_FACE_NAMES.append(facename)
        cap.release()
        cv2.destroyAllWindows()
        print(f"{facename}的人脸录入成功")
    except Exception as ex:
        raise ex


def check_face_reco_handle(num, filename):
    # 人脸识别照片并返回名称
    name = ""
    try:
        global KNOW_FACE_DIR
        global KNOW_FACE_ENCODING
        global KNOW_FACE_NAMES
    except:
        print("请在主程序中导入人脸识别初始化模块")
        return name

    WIDTH = 240
    HEIGHT = 320
    print(f"开始人脸识别...")
    try:
        cap = cv2.VideoCapture(num)
        while not cap.isOpened():
            continue
        # cv2.waitKey(1)
        ret, frame = cap.read()
        rgb_frame = frame[:, :, ::-1]

        face_locations = face_recognition.face_locations(rgb_frame)
        face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)

        for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
            matches = face_recognition.compare_faces(KNOW_FACE_ENCODING, face_encoding)

            if True in matches:
                first_match_index = matches.index(True)
                name = KNOW_FACE_NAMES[first_match_index]

        cv2.imwrite(f'./{filename}', frame)

        cap.release()
        cv2.destroyAllWindows()
        print(f"人脸识别完成: {name}")
        return name
    except Exception as ex:
        raise ex

